from core.data_loader import DataLoader
import pandas as pd
import sqlalchemy

class DatabasePlugin:
    """数据库插件"""
    
    def __init__(self, loader: DataLoader):
        self.loader = loader
        self.supported_dbs = ['mysql', 'postgresql']
    
    def register(self):
        """注册插件"""
        for db_type in self.supported_dbs:
            self.loader.supported_sources[db_type] = self._load_db
    
    def _load_db(self, conn_str: str, **kwargs) -> pd.DataFrame:
        """加载数据库数据"""
        table = kwargs.get('table')
        query = kwargs.get('query')
        
        if not table and not query:
            raise ValueError("必须指定表名或查询语句")
        
        engine = sqlalchemy.create_engine(conn_str)
        
        if query:
            return pd.read_sql(query, engine)
        else:
            return pd.read_sql_table(table, engine)